package org.jordens.jdbcspy.jmx;

import javax.management.*;
import java.io.*;
import java.lang.management.*;
import java.util.*;

/**
 * MBean exposing various JDBC statistics gathered from the JDBCSpyDriver.
 *
 * @author Adam Jordens
 */
public class JDBCSpy implements JDBCSpyMBean
{
    /** JMX bean name reference */
    public static final ObjectName NAME;

    /** Default 'resetInterval' - 60000ms or 60s */
    private static final long DEFAULT_RESET_INTERVAL = 60000L;

    /** Constant for Reset Message */
    private static final String RESET_MSG = "-- RESET --";

    /** All query statistics, SQL -> QueryStatistic */
    private final Map<String, QueryStatistic> executedQueries = new HashMap<String, QueryStatistic>();

    /** Timer used to schedule the periodic reset task */
    private final Timer timer = new Timer();

    /** All statistics are reset every 'resetInterval' # of milliseconds */
    private long resetInterval;

    /** Log writing facility */
    private PrintWriter logFileWriter = null;

    /** Whether or not statistics gathering is enabled. Default is true. */
    private boolean isEnabled = true;

    static
    {
        try
        {
            NAME = new ObjectName(JDBCSpyMBean.NAME);
        }
        catch (Exception e)
        {
            throw new IllegalStateException(e);
        }
    }

    /**
     * Constructs a new JDBCSpy w/ a resetInterval = DEFAULT_RESET_INTERVAL
     */
    public JDBCSpy()
    {
        setResetInterval(DEFAULT_RESET_INTERVAL);
    }

    /**
     * {@inheritDoc}
     */
    public Collection<QueryStatistic> listQueryStatistics()
    {
        List<QueryStatistic> statistics = new ArrayList<QueryStatistic>();
        synchronized(executedQueries)
        {
            statistics.addAll(executedQueries.values());
        }

        Collections.sort(statistics);

        return statistics;
    }

    /**
     * {@inheritDoc}
     */
    public Long getNumberOfQueriesExecuted()
    {
        long count = 0;
        for (QueryStatistic statistic : listQueryStatistics())
        {
            count += statistic.getCount();
        }

        return count;
    }

    /**
     * {@inheritDoc}
     */
    public Long getSQLExecutionTime()
    {
        long time = 0;
        for (QueryStatistic statistic : listQueryStatistics())
        {
            time += statistic.getSQLExecutionTime();
        }

        return time;
    }

    /**
     * {@inheritDoc}
     */
    public synchronized long getResetInterval()
    {
        return resetInterval;
    }

    /**
     * {@inheritDoc}
     */
    public synchronized void setResetInterval(long ms)
    {
        this.resetInterval = ms;

        timer.purge();

        if (ms > 0)
        {
            timer.schedule(new ResetTask(), ms);
        }
    }

    /**
     * Log the given query.
     *
     * This method will do nothing in the event that statistics gathering is
     * not enabled.
     *
     * @param sql SQL executed
     * @param ms Execution time
     */
    public void logQuery(String sql, long ms)
    {
        if (!isEnabled)
        {
            return;
        }

        synchronized(executedQueries)
        {
            QueryStatistic statistic = executedQueries.get(sql);
            if (statistic == null)
            {
                statistic = new QueryStatistic(sql, ms);
            }
            else
            {
                statistic.record(ms);
            }
            executedQueries.put(sql, statistic);
        }

        writeLog(sql, ms + "ms");
    }

    /**
     * {@inheritDoc}
     */
    public void reset()
    {
        synchronized(executedQueries)
        {
            executedQueries.clear();
        }

        writeLog(RESET_MSG);
    }

    /**
     * {@inheritDoc}
     */
    public void enable()
    {
        this.isEnabled = true;
    }

    /**
     * {@inheritDoc}
     */
    public void disable()
    {
        this.isEnabled = false;
    }

    /**
     * Register this MBean on all available servers.
     *
     * If the MBean was previously registered, it will be unregistered and
     * registered again.
     */
    public synchronized void registerBean()
    {
        List<MBeanServer> mbeanServers = MBeanServerFactory.findMBeanServer(null);
        mbeanServers.add(ManagementFactory.getPlatformMBeanServer());

        for (MBeanServer server : mbeanServers)
        {
            try
            {
                server.unregisterMBean(NAME);
            }
            catch (Exception e) { /* do nothing */ }

            try
            {
                server.registerMBean(this, NAME);
            }
            catch (Exception e)
            {
                // TODO-AJ Switch to a proper logging framework
                System.err.println("Unable to register " + NAME + " on " + mbeanServers.get(0));
                e.printStackTrace(System.err);
            }
        }
    }

    /**
     * Specify a file that will be used for logging.
     *
     * In the event that there is already an active log file, the previous file
     * will be automatically closed and a new one created in its place.
     *
     * @param file File that is both readable and writable
     */
    public synchronized void setLogFile(File file)
    {
        if (logFileWriter != null)
        {
            logFileWriter.close();
        }

        try
        {
            logFileWriter = new PrintWriter(new BufferedWriter(new FileWriter(file)));
        }
        catch (IOException e)
        {
            System.err.println("Unable to set log file to " + file);
            e.printStackTrace(System.err);
        }
    }

    /**
     * {@inheritDoc}
     */
    public boolean mark(String message)
    {
        return writeLog(message);
    }

    /**
     * Writes a given message to the current log file.  If not log file is active
     * than this operation does nothing.
     *
     * @param msgs Message(s) to write to log file
     * @return whether or not a message was logged (will return false if logging is disabled)
     */
    private boolean writeLog(Object ... msgs)
    {
        if (logFileWriter != null)
        {
            StringBuilder builder = new StringBuilder();
            builder.append(System.currentTimeMillis());
            builder.append(" - ");
            for (Object msg : msgs)
            {
                builder.append(msg.toString());
                builder.append(" - ");
            }
            logFileWriter.println(builder.toString());
            logFileWriter.flush();

            return true;
        }

        return false;
    }

    /**
     * A basic reset task that will periodically be run to reset all gathered
     * statistics.
     */
    private class ResetTask extends TimerTask
    {
        /**
         * {@inheritDoc}
         */
        public void run()
        {
            reset();
        }
    }
}
